#!/usr/bin/env python3
"""
检查FastAPI WebSocket状态的正确方法
"""

from fastapi import WebSocket
from starlette.websockets import WebSocketState
import logging

logger = logging.getLogger(__name__)

def check_websocket_state(websocket: WebSocket) -> bool:
    """
    检查WebSocket是否已连接
    返回True表示已连接，False表示已断开
    """
    try:
        # FastAPI/Starlette中WebSocket状态的正确检查方法
        if hasattr(websocket, 'client_state'):
            state = websocket.client_state == WebSocketState.CONNECTED
            state_name = str(websocket.client_state).split('.')[-1] if hasattr(websocket.client_state, 'split') else str(websocket.client_state)
            logger.debug(f"WebSocket client_state: {state_name}, connected: {state}")
            return state
        elif hasattr(websocket, 'state'):
            state = websocket.state == WebSocketState.CONNECTED
            state_name = str(websocket.state).split('.')[-1] if hasattr(websocket.state, 'split') else str(websocket.state)
            logger.debug(f"WebSocket state: {state_name}, connected: {state}")
            return state
        else:
            # 获取对象的所有属性，尝试发现状态信息
            try:
                attrs = [attr for attr in dir(websocket) if not attr.startswith('_')]
                state_attrs = {attr: getattr(websocket, attr) for attr in attrs if 'state' in attr.lower()}
                logger.debug(f"WebSocket对象属性: {state_attrs}")
            except Exception as attr_e:
                logger.debug(f"获取WebSocket属性失败: {attr_e}")
            
            # 兜底方案：假设连接正常，实际发送时会捕获异常
            logger.debug("无法检查WebSocket状态，使用兜底方案")
            return True
    except Exception as e:
        logger.debug(f"检查WebSocket状态失败: {e}", exc_info=True)
        return False

async def safe_send_json(websocket: WebSocket, data: dict) -> bool:
    """
    安全发送JSON消息
    """
    try:
        # 先检查连接状态
        if not check_websocket_state(websocket):
            logger.warning("WebSocket已断开，跳过发送消息")
            return False
        
        # 尝试发送消息
        logger.debug(f"尝试发送JSON消息: {data}")
        await websocket.send_json(data)
        logger.debug("JSON消息发送成功")
        return True
    except Exception as e:
        # 区分不同类型的异常并提供详细信息
        error_type = type(e).__name__
        error_msg = str(e)
        
        if "WebSocketDisconnect" in error_type:
            logger.warning(f"WebSocket已断开，无法发送消息: {error_msg}")
        elif "close message has been sent" in error_msg:
            logger.warning(f"WebSocket已关闭，无法发送消息: {error_msg}")
        elif "ConnectionClosed" in error_type:
            logger.warning(f"连接已关闭: {error_msg}")
        else:
            logger.error(f"发送JSON消息失败: {error_type}: {error_msg}", exc_info=True)
            logger.error(f"尝试发送的数据类型: {type(data)}, 数据内容: {data}")
        return False

async def safe_close_websocket(websocket: WebSocket, code: int = 1000, reason: str = "Normal closure") -> bool:
    """
    安全关闭WebSocket连接
    """
    try:
        # 检查连接状态，只有在连接时才尝试关闭
        if check_websocket_state(websocket):
            try:
                await websocket.close(code=code, reason=reason)
                logger.info(f"WebSocket连接已关闭: {code} - {reason}")
                return True
            except RuntimeError as rt_error:
                # 如果是"已经发送了关闭消息"的错误，则认为是正常的
                if "close message has been sent" in str(rt_error):
                    logger.info(f"WebSocket已经处于关闭状态: {rt_error}")
                    return True
                else:
                    # 其他RuntimeError重新抛出
                    raise
        else:
            logger.info("WebSocket已经断开，无需关闭")
            return True
    except Exception as e:
        # 记录详细的异常信息，但不作为错误处理
        logger.debug(f"关闭WebSocket连接时发生异常: {e}", exc_info=True)
        return False
